

using namespace AscendC;

//12.917339
//12.136048
//10.156779

template <typename T>
class TokenShiftB1T1Align_1Core
{
public:
    __aicore__ inline TokenShiftB1T1Align_1Core() {}
    __aicore__ inline void Init(
        GM_ADDR x, 
        GM_ADDR rwkvag, 
        GM_ADDR h0, 
        GM_ADDR xr, 
        GM_ADDR xw, 
        GM_ADDR xk,
        GM_ADDR xv, 
        GM_ADDR xa, 
        GM_ADDR xg, 
        GM_ADDR ht, 
        GM_ADDR workspace,
        uint32_t aivNum,
        uint32_t B_length,
        uint32_t T_length,
        uint32_t C_length,
        uint32_t tiling_size,
        uint32_t core_rows,
        uint32_t core_rows_remain,
        uint32_t X_length_align,
        TPipe* pipeIn)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero");
        this->pipe = pipeIn;
        this->aivNum = aivNum;

        this->B_length = B_length;
        this->T_length = T_length;
        this->C_length = C_length;
        this->tiling_size = tiling_size;
        this->ALIGN_NUM = 32 / sizeof(T);
        uint32_t core_length = this->core_rows * C_length;

        x_global.SetGlobalBuffer((__gm__ T *)x, core_length);

        rwkvag_global.SetGlobalBuffer((__gm__ T *)rwkvag, 6 * C_length);
        h0_global.SetGlobalBuffer((__gm__ T *)h0 , C_length);
        
        xr_global.SetGlobalBuffer((__gm__ T *)xr, core_length);   
        xw_global.SetGlobalBuffer((__gm__ T *)xw, core_length);   
        xk_global.SetGlobalBuffer((__gm__ T *)xk, core_length);   
        xv_global.SetGlobalBuffer((__gm__ T *)xv, core_length);   
        xa_global.SetGlobalBuffer((__gm__ T *)xa, core_length);   
        xg_global.SetGlobalBuffer((__gm__ T *)xg, core_length);        
        ht_global.SetGlobalBuffer((__gm__ T *)ht, C_length); 
        pipe->InitBuffer(inQueueXBind, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(inQueueX, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(inQueueRwkvag, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(inQueueH0, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(inQueueXNextLine, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(outQueueXR, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(outQueueXW, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(outQueueXK, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(outQueueXV, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(outQueueXA, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(outQueueXG, BUFFER_NUM, this->tiling_size * sizeof(T));

        pipe->InitBuffer(QueueTmpX, this->tiling_size * sizeof(T));
        pipe->InitBuffer(QueueTmp0, this->tiling_size * sizeof(T));
        pipe->InitBuffer(QueueTmp1, this->tiling_size * sizeof(T));
        pipe->InitBuffer(QueueTmp2, this->tiling_size * sizeof(T));
        pipe->InitBuffer(QueueTmp3, this->tiling_size * sizeof(T));
        pipe->InitBuffer(QueueTmp4, this->tiling_size * sizeof(T));
        pipe->InitBuffer(QueueTmp5, this->tiling_size * sizeof(T));


        g_RWKVAG0Local = QueueTmp0.Get<T>();
        g_RWKVAG1Local = QueueTmp1.Get<T>();
        g_RWKVAG2Local = QueueTmp2.Get<T>();
        g_RWKVAG3Local = QueueTmp3.Get<T>();
        g_RWKVAG4Local = QueueTmp4.Get<T>();
        g_RWKVAG5Local = QueueTmp5.Get<T>();

        DataCopy(g_RWKVAG0Local, rwkvag_global, C_length);   
        DataCopy(g_RWKVAG1Local, rwkvag_global[C_length], C_length); 
        DataCopy(g_RWKVAG2Local, rwkvag_global[C_length * 2], C_length);
        DataCopy(g_RWKVAG3Local, rwkvag_global[C_length * 3], C_length);
        DataCopy(g_RWKVAG4Local, rwkvag_global[C_length * 4], C_length);
        DataCopy(g_RWKVAG5Local, rwkvag_global[C_length * 5], C_length);

        PipeBarrier<PIPE_ALL>();

    }

    __aicore__ inline void Process() {

        CopyInH0(0, this->tiling_size);
        // CopyInXNextLine(0, this->tiling_size);
        CopyInX(0, this->tiling_size);
        Compute(0, this->tiling_size);
        CopyOutX(0, this->tiling_size);
        // CopyOutHT(0, this->tiling_size);
               
    }
    
    __aicore__ inline void Compute(uint32_t startpos, uint32_t length) {
        LocalTensor<T> xLocal = inQueueXBind.DeQue<T>();
        LocalTensor<T> h0Local = inQueueH0.DeQue<T>();

        LocalTensor<T> XRLocal = outQueueXR.AllocTensor<T>();
        LocalTensor<T> XWLocal = outQueueXW.AllocTensor<T>();
        LocalTensor<T> XKLocal = outQueueXK.AllocTensor<T>();
        LocalTensor<T> XVLocal = outQueueXV.AllocTensor<T>();
        LocalTensor<T> XALocal = outQueueXA.AllocTensor<T>();
        LocalTensor<T> XGLocal = outQueueXG.AllocTensor<T>();

        Sub(h0Local, h0Local, xLocal, length);
        Mul(XRLocal, h0Local, g_RWKVAG0Local, length);
        Add(XRLocal, XRLocal, xLocal, length);
        outQueueXR.EnQue(XRLocal);  

        Mul(XWLocal, h0Local, g_RWKVAG1Local, length);
        Add(XWLocal, XWLocal, xLocal, length);
        outQueueXW.EnQue(XWLocal);  
        

        Mul(XKLocal, h0Local, g_RWKVAG2Local, length);
        Add(XKLocal, XKLocal, xLocal, length);
        outQueueXK.EnQue(XKLocal);  

        Mul(XVLocal, h0Local, g_RWKVAG3Local, length);
        Add(XVLocal, XVLocal, xLocal, length);
        outQueueXV.EnQue(XVLocal); 

        Mul(XALocal, h0Local, g_RWKVAG4Local, length);
        Add(XALocal, XALocal, xLocal, length);
        outQueueXA.EnQue(XALocal); 


        Mul(XGLocal, h0Local, g_RWKVAG5Local, length);
        Add(XGLocal, XGLocal, xLocal, length);
        outQueueXG.EnQue(XGLocal);      
        
        // Adds(h0Local, xLocal, T(0), length);
        DataCopy(ht_global, xLocal, length); 

        inQueueXBind.FreeTensor(xLocal);
    }    

    __aicore__ inline void CopyInH0(uint32_t startpos, uint32_t length) {
        LocalTensor<T> h0Local = inQueueH0.AllocTensor<T>();
        DataCopy(h0Local, h0_global[startpos], length);   
        inQueueH0.EnQue(h0Local);         
    }       
   
    __aicore__ inline void CopyInXNextLine(uint32_t startpos, uint32_t length) {
        LocalTensor<T> xLocal = inQueueXNextLine.AllocTensor<T>();
        DataCopy(xLocal, x_global[startpos], length);   
        inQueueXNextLine.EnQue(xLocal);  
    }
    __aicore__ inline void CopyInX(uint32_t startpos, uint32_t length) {
        LocalTensor<T> xLocal = inQueueXBind.AllocTensor<T>();
        DataCopy(xLocal, x_global[startpos], length);   
        inQueueXBind.EnQue(xLocal);  
    }    

    // __aicore__ inline void CopyInRline(uint32_t startpos, uint32_t length) {
    //     LocalTensor<T> rwkvagLocal = inQueueRwkvag.AllocTensor<T>();
    //     DataCopy(rwkvagLocal, rwkvag_global[startpos], length);   
    //     inQueueRwkvag.EnQue(rwkvagLocal);  
    // }

    // __aicore__ inline void CopyOutHT(uint32_t startpos, uint32_t length) {
    //     // LocalTensor<T> xLocal = inQueueXBind.DeQue<T>();
    //     // printf("CopyOutHT");
    //     DataCopy(ht_global[startpos], h0Local, length);  
    //     PipeBarrier<PIPE_ALL>();  
    // }

    __aicore__ inline void CopyOutX(uint32_t startpos, uint32_t length) {
        LocalTensor<T> xrLocal = outQueueXR.DeQue<T>();
        DataCopy(xr_global[startpos], xrLocal, length); 
        outQueueXR.FreeTensor(xrLocal);

        LocalTensor<T> xwLocal = outQueueXW.DeQue<T>();
        DataCopy(xw_global[startpos], xwLocal, length); 
        outQueueXW.FreeTensor(xwLocal);

        LocalTensor<T> xkLocal = outQueueXK.DeQue<T>();
        DataCopy(xk_global[startpos], xkLocal, length); 
        outQueueXK.FreeTensor(xkLocal);

        LocalTensor<T> xvLocal = outQueueXV.DeQue<T>();
        DataCopy(xv_global[startpos], xvLocal, length); 
        outQueueXV.FreeTensor(xvLocal);

        LocalTensor<T> xaLocal = outQueueXA.DeQue<T>();
        DataCopy(xa_global[startpos], xaLocal, length); 
        outQueueXA.FreeTensor(xaLocal);

        LocalTensor<T> xgLocal = outQueueXG.DeQue<T>();
        DataCopy(xg_global[startpos], xgLocal, length); 
        outQueueXG.FreeTensor(xgLocal);
        PipeBarrier<PIPE_ALL>();
    }

private:
    GlobalTensor<T> x_global;
    GlobalTensor<T> rwkvag_global;
    GlobalTensor<T> h0_global;
    GlobalTensor<T> xr_global, xw_global, xk_global, xv_global, xa_global, xg_global, ht_global; 

    TQueBind<TPosition::VECIN, TPosition::VECOUT, BUFFER_NUM> inQueueXBind;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX, inQueueRwkvag, inQueueH0;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueXNextLine;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueXR, outQueueXW, outQueueXK, outQueueXV, outQueueXA, outQueueXG;

    TBuf<QuePosition::VECCALC> QueueTmpX, QueueTmp0, QueueTmp1, QueueTmp2, QueueTmp3, QueueTmp4, QueueTmp5;
    uint32_t B_length;
    uint32_t T_length;
    uint32_t C_length;
    uint32_t tiling_size;
    uint32_t core_rows;
    uint32_t core_rows_remain;
    uint32_t startPointer;
    uint32_t X_length_align;
    uint32_t ALIGN_NUM;
    uint32_t prevLinePointer;
    uint32_t aivNum;


    LocalTensor<T> g_RWKVAG0Local;
    LocalTensor<T> g_RWKVAG1Local;
    LocalTensor<T> g_RWKVAG2Local;
    LocalTensor<T> g_RWKVAG3Local;
    LocalTensor<T> g_RWKVAG4Local;
    LocalTensor<T> g_RWKVAG5Local;
    AscendC::TPipe* pipe;
};